import datetime
from tensorflow import keras
import numpy as np
from tensorflow.keras import Input
from tensorflow.keras import layers
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Conv2D, Conv2DTranspose, Flatten, Dropout
from tensorflow.keras.optimizers import Adam, RMSprop
from tensorflow.keras import metrics
from keras.models import Model
import mat_util

def automap_model(img_h=64, img_w=64):
    img_size = img_w * img_h
    input_shape = (img_size * 2,)
    x_input = Input(shape=input_shape)
    x = Dense(img_size, activation='tanh')(x_input)
    x = Dropout(0.5)(x)
    x = Dense(img_size, activation='tanh')(x)
    x = Dropout(0.5)(x)
    x_3d = layers.Reshape((img_h, img_w, 1))(x)
    conv1 = Conv2D(64, (5, 5), padding='same', activation='relu')(x_3d)
    conv2 = Conv2D(64, (5, 5), padding='same', activation='relu')(conv1)
    conv3 = Conv2DTranspose(1, (7, 7), activation="relu", strides=(1, 1), padding="same")(conv2)
    out = Flatten()(conv3)

    model = Model(inputs=x_input, outputs=out)
    adam_opt =Adam(learning_rate=0.0002)
    model.compile(optimizer=adam_opt,
                  loss='mse',
                  metrics=['accuracy'])

    return model


if __name__ == '__main__':
    model = automap_model()
    model.summary()
    # ds_x = load_mat(r'E:/experiment/project\keras-automap-master/keras-automap-master/data/mri_train_64x64_2d_images_c30e50.mat')
    # ds_y = load_mat(r'E:/experiment/project/keras-automap-master/keras-automap-master/data/mri_train_64x64_2d_images_c100e100.mat')
    ds_x = mat_util.load_mat(r'E:\experiment\project\keras-automap-master\keras-automap-master\data\train_input.mat')
    ds_y =  mat_util.load_mat(r'E:\experiment\project\keras-automap-master\keras-automap-master\data\train_x_img.mat')
    ds_x = ds_x.T
    ds_y = ds_y.T
    test_x =  mat_util.load_mat(r'E:\experiment\project\keras-automap-master\keras-automap-master\data\test_input.mat"')

    test_x = test_x.T
    print(f'{ds_x.shape=}')
    print(f'{ds_y.shape=}')
    print(f'{test_x.shape=}')
    model.fit(x=ds_x, y=ds_y, validation_split=0.3, verbose=2, epochs=100, batch_size=32)
    predictions = model.predict(test_x)
    print(f'{predictions.shape=}')
    mat_util.save_mat(r'E:\experiment\project\keras-automap-master\keras-automap-master\data\test_64x64_predicted.mat',predictions.T)


    # 保存模型
    dt = datetime.datetime.now()
    model_save_filepath = r'E:/experiment/project/keras-automap-master/keras-automap-master/data/' + dt.strftime('%Y%m%d-%H%M%S') + '.h5'
    model.save(model_save_filepath)
